package org.pytorch.imagesegmentation;

import androidx.appcompat.app.AppCompatActivity;

import android.content.Context;
import android.content.Intent;
import android.graphics.Bitmap;
import android.graphics.BitmapFactory;
import android.net.Uri;
import android.os.Bundle;
import android.os.SystemClock;
import android.provider.MediaStore;
import android.util.Log;
import android.view.View;
import android.widget.Button;
import android.widget.ImageView;
import android.widget.ProgressBar;

import org.pytorch.IValue;
import org.pytorch.LiteModuleLoader;
import org.pytorch.Module;
import org.pytorch.Tensor;
import org.pytorch.torchvision.TensorImageUtils;

import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.Map;

/**
 * 主活动类，负责处理图像分割的功能
 */
public class MainActivity extends AppCompatActivity implements Runnable {

    // 图像视图，用于显示原始图像和分割后的图像
    private ImageView mImageView;
    // 按钮，用于触发图像分割操作
    private Button mButtonSegment;
    // 进度条，用于显示图像分割的进度
    private ProgressBar mProgressBar;
    // 原始图像的位图对象
    private Bitmap mBitmap = null;
    // 加载的图像分割模型
    private Module mModule = null;
    // 当前显示的图像名称
    private String mImagename = "deeplab.jpg";

    private static final int PICK_IMAGE = 1;
    // 图像分割的类别数量
    private static final int CLASSNUM = 21;
    // 狗的类别索引
    private static final int DOG = 12;
    // 人的类别索引
    private static final int PERSON = 15;
    // 羊的类别索引
    private static final int SHEEP = 17;

//    /**
//     * 获取资产文件的路径
//     *
//     * @param context 上下文对象
//     * @param assetName 资产文件的名称
//     * @return 资产文件的路径
//     * @throws IOException 如果文件不存在或读取失败
//     */

    // 处理用户选择的图片
    @Override
    protected void onActivityResult(int requestCode, int resultCode, Intent data) {
        super.onActivityResult(requestCode, resultCode, data);
        if (requestCode == PICK_IMAGE && resultCode == RESULT_OK && data != null) {
            Uri selectedImageUri = data.getData();
            try {
                // 将图片转换为 Bitmap
                mBitmap = BitmapFactory.decodeStream(getContentResolver().openInputStream(selectedImageUri));
                mImageView.setImageBitmap(mBitmap);  // 显示选中的图片
            } catch (FileNotFoundException e) {
                e.printStackTrace();
            }
        }
    }

    public static String assetFilePath(Context context, String assetName) throws IOException {
        // 在应用的文件目录下创建一个新的文件对象
        File file = new File(context.getFilesDir(), assetName);
        // 如果文件已经存在且长度大于0，则直接返回文件的绝对路径
        if (file.exists() && file.length() > 0) {
            return file.getAbsolutePath();
        }

        // 打开资产文件的输入流
        try (InputStream is = context.getAssets().open(assetName)) {
            // 创建一个文件输出流，用于将资产文件写入到新创建的文件中
            try (OutputStream os = new FileOutputStream(file)) {
                // 创建一个缓冲区，用于读取和写入数据
                byte[] buffer = new byte[4 * 1024];
                int read;
                // 循环读取输入流中的数据，并将其写入到输出流中
                while ((read = is.read(buffer)) != -1) {
                    os.write(buffer, 0, read);
                }
                // 刷新输出流，确保所有数据都被写入到文件中
                os.flush();
            }
            // 返回文件的绝对路径
            return file.getAbsolutePath();
        }
    }



        // 获取选择图片按钮


    @Override
    protected void onCreate(Bundle savedInstanceState) {
        super.onCreate(savedInstanceState);
        setContentView(R.layout.activity_main);

        try {
            // 从资产文件中读取图像，并将其解码为位图对象
            mBitmap = BitmapFactory.decodeStream(getAssets().open(mImagename));
        } catch (IOException e) {
            // 如果读取失败，记录错误日志并结束活动
            Log.e("ImageSegmentation", "Error reading assets", e);
            finish();
        }

        // 获取图像视图的引用，并将原始图像设置为其显示内容
        mImageView = findViewById(R.id.imageView);
        mImageView.setImageBitmap(mBitmap);

        // 获取重启按钮的引用，并设置其点击事件监听器
        final Button buttonRestart = findViewById(R.id.restartButton);
        buttonRestart.setOnClickListener(new View.OnClickListener() {
            public void onClick(View v) {
                // 切换显示的图像名称
                if (mImagename == "deeplab.jpg")
                    mImagename = "dog.jpg";
                else
                    mImagename = "deeplab.jpg";
                try {
                    // 从资产文件中读取新的图像，并将其解码为位图对象
                    mBitmap = BitmapFactory.decodeStream(getAssets().open(mImagename));
                    // 将新的图像设置为图像视图的显示内容
                    mImageView.setImageBitmap(mBitmap);
                } catch (IOException e) {
                    // 如果读取失败，记录错误日志并结束活动
                    Log.e("ImageSegmentation", "Error reading assets", e);
                    finish();
                }
            }
        });

        // 获取分割按钮的引用，并设置其点击事件监听器
        mButtonSegment = findViewById(R.id.segmentButton);
        mProgressBar = (ProgressBar) findViewById(R.id.progressBar);
        mButtonSegment.setOnClickListener(new View.OnClickListener() {
            public void onClick(View v) {
                // 禁用分割按钮，显示进度条，并更新按钮文本
                mButtonSegment.setEnabled(false);
                mProgressBar.setVisibility(ProgressBar.VISIBLE);
                mButtonSegment.setText(getString(R.string.run_model));

                // 创建一个新的线程，并启动它来运行图像分割操作
                Thread thread = new Thread(MainActivity.this);
                thread.start();
            }
        });
        Button selectImageButton = findViewById(R.id.selectImageButton);
        selectImageButton.setOnClickListener(new View.OnClickListener() {
            @Override
            public void onClick(View v) {
                // 打开相册选择图片
                Intent intent = new Intent(Intent.ACTION_PICK, MediaStore.Images.Media.EXTERNAL_CONTENT_URI);
                startActivityForResult(intent, PICK_IMAGE);
            }
        });
        try {
            // 加载预训练的图像分割模型
            mModule = LiteModuleLoader.load(MainActivity.assetFilePath(getApplicationContext(), "deeplabv3_scripted_optimized.ptl"));
        } catch (IOException e) {
            // 如果加载失败，记录错误日志并结束活动
            Log.e("ImageSegmentation", "Error reading assets", e);
            finish();
        }

    }

    @Override
    public void run() {
        // 将原始图像转换为张量，并进行归一化处理
        final Tensor inputTensor = TensorImageUtils.bitmapToFloat32Tensor(mBitmap,
                TensorImageUtils.TORCHVISION_NORM_MEAN_RGB, TensorImageUtils.TORCHVISION_NORM_STD_RGB);
        // 获取张量的数据数组
        final float[] inputs = inputTensor.getDataAsFloatArray();

        // 记录开始时间
        final long startTime = SystemClock.elapsedRealtime();
        // 使用加载的模型进行推理，得到输出张量
        Map<String, IValue> outTensors = mModule.forward(IValue.from(inputTensor)).toDictStringKey();
        // 计算推理时间
        final long inferenceTime = SystemClock.elapsedRealtime() - startTime;
        // 记录推理时间 1
        Log.d("ImageSegmentation",  "inference time (ms): " + inferenceTime);

        // 获取输出张量中的 "out" 张量
        final Tensor outputTensor = outTensors.get("out").toTensor();
        // 获取输出张量的数据数组
        final float[] scores = outputTensor.getDataAsFloatArray();
        // 获取原始图像的宽度和高度
        int width = mBitmap.getWidth();
        int height = mBitmap.getHeight();
        // 创建一个整数数组，用于存储分类结果
        int[] intValues = new int[width * height];
        // 遍历每个像素，将其分类为不同的类别
        for (int j = 0; j < height; j++) {
            for (int k = 0; k < width; k++) {
                int maxi = 0, maxj = 0, maxk = 0;
                double maxnum = -Double.MAX_VALUE;
                for (int i = 0; i < CLASSNUM; i++) {
                    // 获取当前像素在每个类别上的得分
                    float score = scores[i * (width * height) + j * width + k];
                    // 找到得分最高的类别
                    if (score > maxnum) {
                        maxnum = score;
                        maxi = i; maxj = j; maxk = k;
                    }
                }
                // 根据类别索引，将像素设置为不同的颜色
                if (maxi == PERSON)
                    intValues[maxj * width + maxk] = 0xFFFF0000;
                else if (maxi == DOG)
                    intValues   [maxj * width + maxk] = 0xFF00FF00;
                else if (maxi == SHEEP)
                    intValues[maxj * width + maxk] = 0xFF0000FF;
                else
                    intValues[maxj * width + maxk] = 0xFF000000;
            }
        }

        // 创建一个新的位图对象，用于存储分类结果
        Bitmap bmpSegmentation = Bitmap.createScaledBitmap(mBitmap, width, height, true);
        Bitmap outputBitmap = bmpSegmentation.copy(bmpSegmentation.getConfig(), true);
        // 将分类结果设置到位图对象中
        outputBitmap.setPixels(intValues, 0, outputBitmap.getWidth(), 0, 0, outputBitmap.getWidth(), outputBitmap.getHeight());
        final Bitmap transferredBitmap = Bitmap.createScaledBitmap(outputBitmap, mBitmap.getWidth(), mBitmap.getHeight(), true);
        // 在主线程中更新UI，显示分割后的图像
        runOnUiThread(new Runnable() {
            @Override
            public void run() {
                mImageView.setImageBitmap(transferredBitmap);
                // 隐藏进度条，启用分割按钮，并更新按钮文本
                mButtonSegment.setEnabled(true);
                mButtonSegment.setText(getString(R.string.segment));
                mProgressBar.setVisibility(ProgressBar.INVISIBLE);

            }
        });
    }
}